import csv
import io

import numpy as np
from fsspec import AbstractFileSystem
from zkl_aiutils_datasets import ChunkedBytesDataset, DelegatingDataset, MappedDataset
from zkl_pyutils_fsspec import FsLike, resolve_fs


def _decode_numpy(sample: bytes) -> np.ndarray:
    buffer = io.BytesIO(sample)
    return np.load(buffer)


class StrEmbeddings(DelegatingDataset[np.ndarray]):
    def __init__(self, fs: FsLike):
        fs = resolve_fs(fs)
        dataset = ChunkedBytesDataset(fs)
        dataset = MappedDataset(dataset, _decode_numpy)
        super().__init__(dataset)
        self._columns_name = load_columns_name(fs)

    @property
    def columns_name(self) -> tuple[str, ...] | None:
        return self._columns_name


def load_columns_name(fs: AbstractFileSystem) -> tuple[str, ...] | None:
    try:
        with fs.open("columns.csv", "rt") as fp:
            reader = csv.reader(fp)
            return tuple(column for column, in reader)
    except FileNotFoundError:
        return None


load = StrEmbeddings
